import json
import sqlite3

import numpy as np
import requests
import pandas as pd
import time
import datetime
import re
from lxml import etree
import efinance as ef
import akshare as ak
from matplotlib import pyplot as plt
import matplotlib.dates as mdates

def parseJson(code):
    with open(f'rawData/{code}.txt') as f:
        json_data = json.load(f)
    # 将 JSON 数据解析为 Python 字典
    # data_dict = json.loads(json_data)

    # 提取 "data" 键对应的值
    data = json_data["data"]

    # 创建 DataFrame
    df = pd.DataFrame(data)
    df["date"] = pd.to_datetime(df["date"])  # 将日期列转换为 datetime 类型

    # 将 "date" 列设置为索引
    df.set_index("date", inplace=True)
    df.sort_index(inplace=True)
    print(df)
    df.to_csv(f'data/{code}.csv')


def parseFundRawData(code):
    df = pd.read_csv(f'rawData/{code}_raw.csv', index_col=0, parse_dates=True)
    dfData = df['totalNetValue']
    dfData.sort_index(inplace=True)
    dfData = dfData.fillna(method='ffill')
    print(dfData)
    dfData.to_csv(f'data/{code}.csv')


# 基金
def download_fund_history_data(fund_code, startDate, endDate, isUpdate):
    url = "https://fundf10.eastmoney.com/F10DataApi.aspx?type=lsjz&code={}&page=1&sdate={}&edate={}&per=40".format(
        fund_code,
        startDate,
        endDate)
    print(url)
    response = requests.get(url)
    print(response.text)
    pages = re.findall('pages:(.*),', response.text)
    int_pages = 0
    if len(pages) > 0:
        int_pages = int(pages[0])
    print("总共{}页".format(int_pages))
    if isUpdate:
        int_pages = 1
    td_all = []
    for i in range(int_pages):
        url = "https://fundf10.eastmoney.com/F10DataApi.aspx?type=lsjz&code={}&page={}&sdate={}&edate={}&per=40".format(
            fund_code,
            i + 1,
            startDate,
            endDate)
        print(url)
        response = requests.get(url)
        time.sleep(0.5)
        print(response.text)
        tables = re.findall('"(.*)"', response.text)
        if len(tables) > 0:
            table = tables[0]
            print(table)
            tree = etree.HTML(table)
            td_list = tree.xpath('//tr/td')
            # print(td_list)
            for t in td_list:
                # print(t.text)
                if t.text == None:
                    td_all = td_all + [""]
                else:
                    td_all.append(t.text)

    net_date = []
    net_a = []
    net_b = []
    net_c = []
    net_d = []
    net_e = []
    for k, v in enumerate(td_all):
        i = k % 7
        if i == 0:
            net_date.append(v)
        if i == 1:
            net_a.append(v)
        if i == 2:
            net_b.append(v)
        if i == 3:
            net_c.append(v)
        if i == 4:
            net_d.append(v)
        if i == 5:
            net_e.append(v)

    df = pd.DataFrame({'date': net_date,
                       'netValue': net_a,
                       'totalNetValue': net_b,
                       'rate': net_c, })
    conn = sqlite3.connect('fund_data.db')
    cursor = conn.cursor()

    # 创建一个表（如果表不存在）
    cursor.execute(f'''
    CREATE TABLE IF NOT EXISTS [{fund_code}] (
        date TEXT PRIMARY KEY,
        netValue INTEGER NOT NULL,
        totalNetValue INTEGER NOT NULL,
        rate INTEGER NOT NULL
    )
    ''')

    # 使用 INSERT OR REPLACE 插入数据
    for index, row in df.iterrows():
        cursor.execute(f'''
        INSERT OR REPLACE INTO [{fund_code}] (date, netValue, totalNetValue,rate) VALUES (?, ?, ?, ?)
        ''', (row['date'], row['netValue'], row['totalNetValue'], row['rate']))

    # 提交事务并关闭连接
    conn.commit()
    conn.close()
    # df.to_sql(fund_code, conn, if_exists='append', dtype={'date': 'TEXT PRIMARY KEY'},index=True, index_label='date')
    # conn.close()
    df.to_csv(f"rawData/{fund_code}_raw.csv".format(fund_code), index=False, encoding="utf-8")


def parseFundDatabase(code):
    conn = sqlite3.connect('fund_data.db')
    query = f"SELECT * FROM [{code}]"
    df = pd.read_sql_query(query, conn)
    conn.close()
    dfData = df[['date', 'totalNetValue']]
    dfData = dfData.set_index('date')
    dfData.sort_index(inplace=True)
    dfData = dfData.fillna(method='ffill')
    print(dfData)
    dfData.to_csv(f'data/{code}.csv')


def test():
    # 获取当前日期
    current_date = datetime.datetime.now().date()
    # 计算前两天的日期
    two_days_ago = current_date - datetime.timedelta(days=2)
    print("当前日期:", datetime.datetime.now().strftime("%Y%m%d"))
    print("前两天的日期:", two_days_ago.strftime("%Y%m%d"))


def updateData(codeList):
    #codeList = ["588000", "512890", "159939", "512690", "515170", "510880"]
    # for code in codeList:
    #     current_date = datetime.datetime.now().date()
    #     days_ago = current_date - datetime.timedelta(days=50)
    #     startDate = days_ago.strftime("%Y%m%d")
    #     endDate = datetime.datetime.now().strftime("%Y%m%d")
    #     download_fund_history_data(code, startDate, endDate, True)
    #     parseFundDatabase(code)
    #     time.sleep(1)
    current_date = datetime.datetime.now().date()
    days_ago = current_date - datetime.timedelta(days=30)
    startDate = days_ago.strftime("%Y%m%d")
    endDate = datetime.datetime.now().strftime("%Y%m%d")
    downloadAllData(codeList,startDate,endDate)

def downloadAllData(codeList,startDate,endDate):
    conn = sqlite3.connect('fund_data.db')
    for code in codeList:
        time.sleep(1)
        df = ef.stock.get_quote_history(stock_codes=code, beg=startDate, end=endDate)
        print(df.head())
        cursor = conn.cursor()
        # 创建一个表（如果表不存在）
        cursor.execute(f'''
               CREATE TABLE IF NOT EXISTS fund (
                   code TEXT NOT NULL,
                   name TEXT NOT NULL,
                   date TEXT NOT NULL,
                   close_price INTEGER NOT NULL,
                   PRIMARY KEY (code, date)
               )
               ''')

        # 使用 INSERT OR REPLACE 插入数据
        for index, row in df.iterrows():
            cursor.execute(f'''
                INSERT OR REPLACE INTO fund (code, name, date,close_price) VALUES (?, ?, ?, ?)
                ''', (row['股票代码'], row['股票名称'], row['日期'], row['收盘']))

        # 提交事务并关闭连接
        conn.commit()
        query = f"SELECT * FROM fund where code={code}"
        df = pd.read_sql_query(query, conn)
        dfData = df[['code','date', 'close_price']]
        dfData = dfData.set_index('date')
        dfData.sort_index(inplace=True)
        dfData = dfData.fillna(method='ffill')
        dfData.to_csv(f'data/{code}.csv')
    conn.close()


def rolling_annualized_return(data, window_years=3):
    """
    计算滚动年化收益率
    参数：
    data : DataFrame，必须包含date和nav两列
    window_years : 滚动窗口年数
    """
    # 预处理数据
    df = data.copy()
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date').set_index('date')

    # 计算滚动窗口天数（假设每年252个交易日）
    window_days = window_years * 252

    # 计算滚动收益率
    rolling_return = (df['close_price'].pct_change(periods=window_days) + 1)

    # 年化处理：(1 + total_return)^(1/years) - 1
    annualized_return = (rolling_return ** (1 / window_years) - 1) * 100

    return annualized_return.dropna()

def judgeNiushiData():
    # 中证偏股型基金指数(CSI:930950)
    stock_zh_index_hist_csindex_df = ak.stock_zh_index_hist_csindex(symbol="930950", start_date="20160101",
                                                                    end_date="20250223")
    print(stock_zh_index_hist_csindex_df)
    df = stock_zh_index_hist_csindex_df[['日期','指数代码','收盘']]
    df = df.rename(columns={'日期': 'date', '指数代码': "code",'收盘':'close_price'})
    df.to_csv("data/930950.csv")

def judgeNiushi():
    # 生成示例数据（2005-2023年沪深300指数）
    df = pd.read_csv("data/930950.csv")

    # 计算3年滚动年化收益
    result = rolling_annualized_return(df, window_years=3)

    # 绘制结果
    plt.figure(figsize=(12, 6))
    plt.plot(result.index, result, label='3年滚动年化收益率', color='#2c7fb8')

    # 格式设置
    plt.title('3年滚动年化收益率走势', fontsize=14, pad=20)
    plt.ylabel('年化收益率 (%)', fontsize=12)
    plt.grid(True)
    plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统
    # plt.rcParams['font.sans-serif'] = ['PingFang HK']  # MacOS
    # plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # Linux
    plt.rcParams['axes.unicode_minus'] = False

    plt.legend()
    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    codeList = ["588000", "512890", "159939", "512690", "515170", "510880","159915","510300","512100","159941","518880","159755","562500","159995","159992","516160","515220"]
    #codeList = ["515220"]
    # parseJson("588000")
    current_date = datetime.datetime.now().date()
    days_ago = current_date - datetime.timedelta(days=3000)
    startDate = days_ago.strftime("%Y%m%d")
    endDate = datetime.datetime.now().strftime("%Y%m%d")

    updateData(codeList)
    #downloadAllData(codeList,startDate,endDate)


